#!/usr/bin/env python3

import message_filters
import numpy as np
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image, PointCloud2
from vision_msgs.msg import Detection2D, Detection2DArray

from .camera_processor import CameraProcessor
from .detector import Detector
from .lidar_processor import LidarProcessor


class FusionNode(Node):
    def __init__(self) -> None:
        super().__init__("fusion_node")

        # 加载配置
        self.load_parameters()

        # 初始化处理器
        self.detector = Detector(self.model_path, self.conf_thres)
        self.lidar_proc = LidarProcessor(self.point_cloud_skip)
        self.camera_proc = CameraProcessor()

        # 创建订阅者
        self.image_sub = message_filters.Subscriber(self, Image, "/camera/image_raw")
        self.lidar_sub = message_filters.Subscriber(self, PointCloud2, "/lidar/points")

        # 同步订阅器
        self.ts = message_filters.ApproximateTimeSynchronizer(
            [self.image_sub, self.lidar_sub],
            10,
            0.1,
        )
        self.ts.registerCallback(self.fusion_callback)

        # 创建发布者
        self.detection_pub = self.create_publisher(
            Detection2DArray,
            "/fusion/detections",
            10,
        )
        self.vis_pub = self.create_publisher(Image, "/fusion/visualization", 10)

        self.get_logger().info("Fusion node has been started")

    def load_parameters(self) -> None:
        """加载参数."""
        self.declare_parameter("model_path", "models/yolov8n.pt")
        self.declare_parameter("conf_thres", 0.5)
        self.declare_parameter("point_cloud_skip", 2)

        self.model_path = self.get_parameter("model_path").value
        self.conf_thres = self.get_parameter("conf_thres").value
        self.point_cloud_skip = self.get_parameter("point_cloud_skip").value

    def fusion_callback(self, img_msg, pc_msg) -> None:
        """处理同步的图像和点云数据."""
        try:
            # 处理图像
            cv_image = self.camera_proc.process_image(img_msg)
            if cv_image is None:
                return

            # 目标检测
            detections = self.detector.detect(cv_image)

            # 处理点云
            points_3d = self.lidar_proc.process_pointcloud(pc_msg)
            points_2d = self.lidar_proc.project_to_image(
                points_3d,
                self.camera_matrix,
                self.lidar_to_camera,
            )

            # 融合处理
            results = self.process_fusion(detections, points_3d, points_2d)

            # 发布结果
            self.publish_results(results, img_msg.header, cv_image)

        except Exception as e:
            self.get_logger().error(f"Error in fusion callback: {e!s}")

    def process_fusion(self, detections, points_3d, points_2d) -> list:
        """处理融合结果."""
        results = []
        for det in detections:
            bbox = det["bbox"]

            # 找到落在检测框内的点云
            mask = (
                (points_2d[:, 0] >= bbox[0])
                & (points_2d[:, 0] <= bbox[2])
                & (points_2d[:, 1] >= bbox[1])
                & (points_2d[:, 1] <= bbox[3])
            )

            points_in_box = points_3d[mask]

            if len(points_in_box) > 0:
                distance = np.linalg.norm(points_in_box.mean(axis=0))
                det["distance"] = distance
                results.append(det)

        return results

    def publish_results(self, results, header, cv_image) -> None:  # 添加cv_image参数
        """发布处理结果."""
        # 发布检测消息
        detection_array = Detection2DArray()
        detection_array.header = header

        for result in results:
            detection = Detection2D()
            # 填充检测信息
            detection_array.detections.append(detection)

        self.detection_pub.publish(detection_array)

        # 发布可视化结果
        if len(results) > 0:
            distances = [r.get("distance", 0) for r in results]
            vis_image = self.camera_proc.draw_detections(
                cv_image.copy(),
                results,
                distances,
            )  # 使用传入的cv_image
            self.vis_pub.publish(
                self.camera_proc.bridge.cv2_to_imgmsg(vis_image, "bgr8"),
            )


def main() -> None:
    rclpy.init()
    node = FusionNode()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()


if __name__ == "__main__":
    main()
